#计算每个物种的基因个数和基因组长度生成_mRNA_countsLength.txt文件
import os
import subprocess
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--path", required=True)
args = parser.parse_args()
dir_path = args.path + "/"

files = os.listdir(dir_path)
dir_path = [os.path.join(dir_path,x).replace("\\","/") for x in files]
species_path = zip(files,dir_path)


for species, path in species_path:
    bed_file = path + "/" + species + ".bed"
    chrom_sizes_file = "/data/genome/IMP/" + species + "/result/igv/" + species + ".chrom.sizes.txt"
    out_path = path + "/" + species + "_mRNA_countsLength.txt"
    if os.path.exists(chrom_sizes_file):
        #print(bed_file,chrom_sizes_file)
        command1 = f"awk -F '\t' '{{print $5}}' '{bed_file}' | sort | uniq | wc -l"    #取bed文件第五列去重得到的个数记为基因个数
        result1 = subprocess.run(command1,shell=True,capture_output=True,text=True)
        gene_num = result1.stdout.strip()
        command2 = f"awk -F '\t' '{{sum+=$2}}END{{print sum}}' '{chrom_sizes_file}'"  #取chrom.sizes.txt文件第二列总和记为基因组长度
        result2 = subprocess.run(command2,shell=True,capture_output=True,text=True)
        genome_num = result2.stdout.strip()
        #print(gene_num,genome_num)
        file = species + "\t" + gene_num + "\t" + genome_num + "\n"
        with open(out_path, 'w') as f:
            f.write(file)
        #file = species + "\t" + gene_num + "\t" + genome_num
        #print(file)
        #print(out_path)

